use dashmap::DashMap;
use futures_util::FutureExt;
use once_cell::sync::Lazy;
use pingora::lb::selection::RoundRobin;
use pingora::lb::{discovery, Backends, LoadBalancer};
use pingora::prelude::TcpHealthCheck;
use pingora::protocols::l4::socket::SocketAddr;
use pingora::{lb::Backend, upstreams::peer::PeerOptions};
use std::time::Duration;
use std::{collections::BTreeSet, net::ToSocketAddrs, sync::Arc};

pub fn new_internal_error(status: u16, message: String) -> pingora::BError {
    pingora::Error::because(
        pingora::ErrorType::HTTPStatus(status),
        message,
        pingora::Error::new(pingora::ErrorType::InternalError),
    )
}

// 全局保存的upstream，用于后续流程中获取
static UPSTREAMS: Lazy<DashMap<String, Arc<Upstream>>> = Lazy::new(|| DashMap::new());

// 添加upstream至全局对象中
pub fn add_upstream(name: &str, up: Upstream) {
    UPSTREAMS.insert(name.to_string(), Arc::new(up));
}

// 获取对应的upstream
pub fn get_upstream(name: &str) -> Option<Arc<Upstream>> {
    if let Some(up) = UPSTREAMS.get(&name.to_string()) {
        return Some(up.clone());
    }
    None
}

pub struct Upstream {
    lb: Arc<LoadBalancer<RoundRobin>>,
    // 是否tls连接
    pub tls: bool,
    // 证书对应的sni
    pub sni: String,
    // 连接节点的相关属性（超时等）
    pub options: PeerOptions,
}

impl Upstream {
    pub fn new(addrs: Vec<String>) -> pingora::Result<Upstream> {
        let mut upstreams = BTreeSet::new();
        let mut backends = Vec::new();
        for addr in addrs.iter() {
            // 需要注意，此处如果是域名则会解析为对应的ip
            // 因此如果是域名且ip动态变化的，需要使用另外的服务发现方式
            for item in addr
                .to_socket_addrs()
                .map_err(|e| new_internal_error(500, e.to_string()))?
            {
                let ext = pingora::lb::Extensions::new();
                backends.push(Backend {
                    addr: SocketAddr::Inet(item),
                    weight: 1,
                    ext,
                });
            }
        }
        upstreams.extend(backends);
        // 静态服务发现，直接使用固定的ip列表，无动态更新
        let discovery = discovery::Static::new(upstreams);
        let backends = Backends::new(discovery);
        // 使用round robin的算法获取upstream节点
        let mut lb = LoadBalancer::<RoundRobin>::from_backends(backends);

        // 使用tcp的方式检测端口，建议使用HttpHealthCheck
        let mut check = TcpHealthCheck::new();
        check.peer_template.options.connection_timeout = Some(Duration::from_secs(3));
        lb.set_health_check(check);

        // 服务发现的刷新间隔，因为是static，所以无需设置
        // lb.update_frequency = Duration::from_secs(60);
        // 健康检测的间隔，根据自己的服务选择合适的值
        lb.health_check_frequency = Some(Duration::from_secs(10));
        // 初始化时首先触发一次更新，生成可使用的upstream节点
        lb.update()
            .now_or_never()
            .expect("static should not block")
            .expect("static should not error");
        Ok(Self {
            lb: Arc::new(lb),
            tls: false,
            sni: "".to_string(),
            options: PeerOptions::new(),
        })
    }

    pub fn lb(&self) -> &Arc<LoadBalancer<RoundRobin>> {
        &self.lb
    }
}
